#!/usr/bin/env python
# encoding: utf-8
"""
gridgen.py

Created by Rob Hetland on 2008-01-18.
Copyright (c) 2008 Texas A&M Univsersity. All rights reserved.
"""

def Gridgen(xbry, ybry, beta, shape, focus=None, ul_idx=0, \
            geographic=False, proj=None, \
            nnodes=14, precision=1.0e-12, newton=1, thin=1, \
            checksimplepoly=1, nppe=3, \
            verbose=False, windows=False):
    """
    create grid object from gridgen code by Pavel Sakov.  See:
        http://www.marine.csiro.au/~sakov/
    
    grid = Gridgen(xbry, ybry, beta, shape, focus=None, star=0,
                 nnodes=14, precision=1.0e-12, newton=1, thin=1,
                 checksimplepoly=1, nppe=3, verbose=False)
    
    input:
        xbry, ybry, beta, ul_idx=0:
            the input polygon represents the boundaries of the domain
            to be calculated.  The polygon should be defined counter-
            clockwise (positive).  Beta represent the locations where the
            boundary makes a 90 turn to the left (beta=1) or the right
            (beta=-1).  The index 'ul_idx' is the index of the upper left
            corner of the domain.  There may be more than one posibility,
            but the output grid will not change.  See the gridgen documentation
            for more details (this is the 'starred' entry in the bry file)
        geographic:
            If geographic is True (default is False), the coordinates are
            projected using the Basemap instance proj.  The grid instance
            that is created is a geographic grid, with both lon/lat and
            x/y defined.
        proj:
            Is the Basemap instance used to project the geographic coordinates
            to cartesian coordinates before calling gridgen.  The default is
                    proj = Basemap(projection='merc', lat_ts=0.0)
        shape:
            the number of points in the grid (ny, nx).  When creating a
            ROMS grid, note that Lm = nx-3 and Mm = ny-3.
        focus:
            The focus function returns values on a grid between 0 and 1,
            given inputs as a uniform unit grid.  E.g.,
                    x, y = mgrid[0:1:100j, 0:1:100j]
                    xf, yf = focus(x, y)
            The resulting grid defines where the output will have increased
            resolution.  If focus=None, the output grid will be uniformly
            distributed.
        other keyword arguments:
            follow the parameter definitions in the gridgen prm file.
    
    returns:
        ROMS Grid instance.
    """
    xbry = asarray(xbry)
    ybry = asarray(ybry)
    beta = asarray(beta)
    
    if proj is None:
        proj = Basemap(projection='merc', resolution=None, lat_ts=0.0)
        
    if geographic:
        xbry, ybry = proj(xbry, ybry)
    
    assert beta.sum() == 4.0, 'sum of beta must be 4.0'
    star = ['']*len(xbry)
    star[ul_idx] = '*'
    ny = shape[0]
    nx = shape[1]
    
    # tempnam gives security warning
    warnings.filterwarnings('ignore', category=RuntimeWarning)
    prmname = os.tempnam()    # prm
    bryname = os.tempnam()    # input
    focusname = os.tempnam()  # grid
    rectname = os.tempnam()   # rectangle
    sigmaname = os.tempnam()  # sigmas
    gridname = os.tempnam()   # output
    
    # Write BRY file (input polygon)
    f = open(bryname, 'w')
    f.writelines(["%f %f %d%s\n" % (x, y, b, s) \
                  for x, y, b, s \
                  in zip(xbry, ybry, beta, star)])
    f.close()
    
    # Write FOCUS file
    if focus is not None:
        y, x = mgrid[0:1:ny*1j, 0:1:nx*1j]
        xfocus, yfocus = focus(x, y)
        f = open(focusname, 'w')
        f.writelines(["%f %f\n" % (x, y) \
                      for x, y \
                      in zip(xfocus.flat, yfocus.flat)])
        f.close()
    
    # Write PRM file
    f = open(prmname, 'w')
    f.write('input %s\n' % bryname)
    f.write('output %s\n' % gridname)
    if focus is not None:
        f.write('grid %s\n' % focusname)
    f.write('nx %d\n' % nx)
    f.write('ny %d\n' % ny)
    f.write('precision %g\n' % precision)
    f.write('thin %d\n' % thin)
    f.write('checksimplepoly %d\n' % checksimplepoly)
    f.write('newton %d\n' % newton)
    f.write('sigmas %s\n' % sigmaname)
    f.write('rectangle %s\n' % rectname)
    f.write('nppe %d\n' % nppe)
    f.close()
    
    # Run gridgen
    if verbose:
        verbstr = '-v'
    else:
        verbstr = ''
    os.system('gridgen %s %s' % (verbstr, prmname))
    
    # Read in grid file
    if windows == False:
        xp, yp = pl.load(gridname).T
    else:
        x=[]; y=[]
        gridfile = open(gridname)
        for line in gridfile:
            if '#' in line:
                x.append(nan)
                y.append(nan)
            else:
                data = line.split()
                x.append(float(data[0]))
                y.append(float(data[1]))
        xp = asarray(x)
        yp = asarray(y)
        gridfile.close()
    xp = xp.reshape(ny, nx)
    yp = yp.reshape(ny, nx)
    # remove temporary files
    
    try:
        [os.remove(file) \
         for file \
         in (prmname, bryname, focusname, rectname, sigmaname, gridname)]
    except:
        pass
    
    if any(isnan(xp)) or any(isnan(yp)):
        xp = ma.masked_where(isnan(xp), xp)
        yp = ma.masked_where(isnan(yp), yp)
    
    if geographic:
        if isinstance(xp, ma.MaskedArray): xp=xp.filled(nan)
        if isinstance(yp, ma.MaskedArray): yp=yp.filled(nan)
        lon, lat = proj(xp, yp, inverse=True)
        lon = ma.masked_where(isnan(lon), lon)
        lat = ma.masked_where(isnan(lat), lat)
        return Grid(lon=lon, lat=lat, proj=proj)
    else:
        return Grid(x=xp, y=yp)
    

